import http from 'k6/http';
import { check, sleep, group, fail } from 'k6';
import { Trend, Counter, Rate, Gauge } from 'k6/metrics';
import { SharedArray } from 'k6/data';
import { randomIntBetween, randomString, uuidv4 } from 'https://jslib.k6.io/k6-utils/1.2.0/index.js';

// ------------------------------
// 1. 自定义指标定义
// ------------------------------
// 连接相关指标
const websocketConnectionTime = new Trend('websocket_connection_time');
const connectionSuccessRate = new Rate('connection_success_rate');
const reconnectAttempts = new Counter('reconnect_attempts');

// 直播交互指标
const chatSendSuccessRate = new Rate('chat_send_success_rate');
const heartBeatInterval = new Trend('heartbeat_interval_ms');

// 系统状态指标
const currentOnlineUsers = new Gauge('current_online_users');
const roomJoinSuccessRate = new Rate('room_join_success_rate');

// ------------------------------
// 2. 测试数据加载
// ------------------------------
const testData = new SharedArray('test_data', function () {
    // 直播房间ID列表（模拟不同直播间）
    return [
        { roomId: 'live-1001', name: '热门游戏直播' },
        { roomId: 'live-1002', name: '才艺表演' },
        { roomId: 'live-1003', name: '产品发布会' },
        { roomId: 'live-1004', name: '户外直播' },
        { roomId: 'live-1005', name: '教育讲座' }
    ];
});

// ------------------------------
// 3. 压测配置 - 梯度上升至50万+在线用户
// ------------------------------
export const options = {
    // 分布式压测配置（单节点无法承受50万用户）
    // 需配合--distributed和--masters参数使用
    vus: 5000,  // 单节点虚拟用户数（根据节点性能调整）
    duration: '2h',  // 总测试时长

    // 阶梯式用户增长策略
    stages: [
        { duration: '5m', target: 1000 },    // 初始热身：1000用户
        { duration: '10m', target: 5000 },   // 第一阶段：5000用户
        { duration: '15m', target: 20000 },  // 第二阶段：2万用户
        { duration: '20m', target: 100000 }, // 第三阶段：10万用户
        { duration: '30m', target: 300000 }, // 第四阶段：30万用户
        { duration: '30m', target: 500000 }, // 第五阶段：50万用户
        { duration: '30m', target: 500000 }, // 稳定期：维持50万用户
        { duration: '10m', target: 0 },      // 逐步退出
    ],

    // 关键指标阈值
    thresholds: {
        connection_success_rate: ['rate>0.99'],    // 连接成功率>99%
        room_join_success_rate: ['rate>0.98'],     // 进入直播间成功率>98%
        chat_send_success_rate: ['rate>0.95'],     // 聊天发送成功率>95%
        http_req_failed: ['rate<0.01'],            // HTTP请求失败率<1%
        reconnect_attempts: ['rate<0.005'],        // 重连率<0.5%
    },

    // 提升WebSocket性能
    setupTimeout: '5m',
    teardownTimeout: '5m',
    throw: true,  // 连接错误时抛出异常
};

// ------------------------------
// 4. 辅助函数
// ------------------------------
// 生成随机用户信息
function generateUser() {
    return {
        userId: `user_${uuidv4().substring(0, 8)}`,
        deviceId: `device_${randomString(12)}`,
        appVersion: `v${randomIntBetween(1, 3)}.${randomIntBetween(0, 9)}.${randomIntBetween(0, 99)}`
    };
}

// 建立WebSocket连接
function connectWebSocket(user, roomId) {
    const start = Date.now();
    const wsUrl = `wss://live-api.example.com/ws?roomId=${roomId}&userId=${user.userId}&deviceId=${user.deviceId}`;

    // 建立WebSocket连接
    const response = http.ws(wsUrl, function (socket) {
        // 连接建立成功
        socket.on('open', () => {
            websocketConnectionTime.add(Date.now() - start);
            connectionSuccessRate.add(true);
        });

        // 处理错误
        socket.on('error', (e) => {
            console.error(`WebSocket错误: ${e}`);
            connectionSuccessRate.add(false);
            reconnectAttempts.add(1);
            fail(`WebSocket连接失败: ${e}`);
        });

        // 处理关闭事件
        socket.on('close', (code, reason) => {
            if (code !== 1000) {  // 非正常关闭
                console.log(`连接关闭: ${code}, 原因: ${reason}`);
                reconnectAttempts.add(1);
            }
        });
    });

    return response;
}

// 发送心跳包维持连接
function sendHeartbeat(socket, user) {
    const interval = randomIntBetween(30000, 45000); // 30-45秒随机心跳间隔
    socket.setInterval(() => {
        const start = Date.now();
        socket.send(JSON.stringify({
            type: 'heartbeat',
            userId: user.userId,
            timestamp: Date.now()
        }));
        heartBeatInterval.add(Date.now() - start);
    }, interval);
}

// 随机发送聊天消息
function sendRandomChat(socket, user) {
    if (Math.random() < 0.3) { // 30%概率发送消息
        const messages = [
            '666666',
            '主播太厉害了！',
            '这个直播不错',
            '有人知道背景音乐吗？',
            '刚来，发生了什么？',
            '赞赞赞！',
            randomString(randomIntBetween(2, 10)) // 随机字符串
        ];

        const message = randomItem(messages);
        const start = Date.now();

        socket.send(JSON.stringify({
            type: 'chat',
            userId: user.userId,
            content: message,
            timestamp: Date.now()
        }));

        // 等待消息确认
        socket.on('message', (data) => {
            const res = JSON.parse(data);
            if (res.type === 'chat_ack' && res.content === message) {
                chatSendSuccessRate.add(true);
            }
        });

        // 5秒后未收到确认则视为失败
        socket.setTimeout(() => {
            chatSendSuccessRate.add(false);
        }, 5000);
    }
}

// ------------------------------
// 5. 测试主体逻辑
// ------------------------------
export default function () {
    const user = generateUser();
    const room = randomItem(testData); // 随机选择一个直播间

    // 记录当前在线用户数
    currentOnlineUsers.add(options.vus);

    group(`用户 ${user.userId} 行为流程`, function () {
        // 1. 获取直播间信息（HTTP请求）
        const roomInfoRes = http.get(`https://live-api.example.com/rooms/${room.roomId}`, {
            headers: {
                'User-Agent': `LiveApp/${user.appVersion}`,
                'X-Device-Id': user.deviceId
            }
        });

        check(roomInfoRes, {
            '获取直播间信息成功': (r) => r.status === 200,
        });

        // 2. 加入直播间（WebSocket连接）
        const wsResponse = connectWebSocket(user, room.roomId);

        check(wsResponse, {
            'WebSocket握手成功': (r) => r.status === 101,
        });

        roomJoinSuccessRate.add(wsResponse.status === 101);

        // 3. 维持连接并进行交互
        wsResponse.socket.on('open', () => {
            // 发送加入房间确认
            wsResponse.socket.send(JSON.stringify({
                type: 'join',
                userId: user.userId,
                roomId: room.roomId,
                timestamp: Date.now()
            }));

            // 启动心跳机制
            sendHeartbeat(wsResponse.socket, user);

            // 模拟用户在直播间的行为（持续随机互动）
            const interactionDuration = randomIntBetween(60000, 300000); // 1-5分钟
            const interactionInterval = randomIntBetween(5000, 15000); // 5-15秒间隔

            const interactionTimer = wsResponse.socket.setInterval(() => {
                sendRandomChat(wsResponse.socket, user);
            }, interactionInterval);

            // 随机时间后离开直播间
            wsResponse.socket.setTimeout(() => {
                wsResponse.socket.clearInterval(interactionTimer);
                wsResponse.socket.send(JSON.stringify({
                    type: 'leave',
                    userId: user.userId,
                    timestamp: Date.now()
                }));
                wsResponse.socket.close(1000, '正常离开');
            }, interactionDuration);
        });

        // 等待用户行为完成
        sleep(randomIntBetween(60, 300)); // 1-5分钟
    });
}

// ------------------------------
// 6. 测试总结
// ------------------------------
export function handleSummary(data) {
    return {
        'live_app_summary.json': JSON.stringify({
            metrics: data.metrics,
            thresholds: data.thresholds,
            summary: {
                totalUsers: options.stages[4].target,
                duration: options.duration,
                successRates: {
                    connection: data.metrics.connection_success_rate.values.rate,
                    roomJoin: data.metrics.room_join_success_rate.values.rate,
                    chatSend: data.metrics.chat_send_success_rate.values.rate,
                }
            }
        }, null, 2),
    };
}
